/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package com.snda.lidejia.mongodbtest;

import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.Mongo;
import com.mongodb.WriteConcern;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 *
 * @author lidejia
 */
public class InsertAndQuery {

    private static final Logger LOGGER = Logger.getLogger(InsertAndQuery.class.getName());
    //
    private static int INSERT_THREAD_NUM = 5;
    private static int QUERY_THREAD_NUM = 40;
    private static int ITEM_SIZE = 10;
    //
    private static long QUERY_RANGE = 0L;
    //
    private static AtomicLong CURRENT_ID = new AtomicLong(0);
    private static AtomicLong LAST_SUCCESS_ID = new AtomicLong(-1);
    //
    private static final String HOST = "116.211.21.86";
    //
    private static String buffer;

    /**
     * @param args the command line arguments
     */
    public static void main(String[] args) throws Exception {

        if (args.length == 5) {
            CURRENT_ID.set(Long.parseLong(args[0]));
            LAST_SUCCESS_ID.set(CURRENT_ID.get() - 1);
            QUERY_RANGE = Long.parseLong(args[1]);
            INSERT_THREAD_NUM = Integer.parseInt(args[2]);
            QUERY_THREAD_NUM = Integer.parseInt(args[3]);
            ITEM_SIZE = Integer.parseInt(args[4]);
        }
        System.out.println(String.format("Insert starting ID is %s", CURRENT_ID.get()));
        System.out.println(String.format("Query range is %s", QUERY_RANGE == 0 ? "unlimit" : QUERY_RANGE));
        System.out.println(String.format("Insert thread number is %s", INSERT_THREAD_NUM));
        System.out.println(String.format("Query thread number is %s", QUERY_THREAD_NUM));
        System.out.println(String.format("Item size is %s", ITEM_SIZE));
        //
        System.setProperty("MONGO.POOLSIZE", Integer.toString(QUERY_THREAD_NUM + INSERT_THREAD_NUM));
        //
        StringBuilder sb = new StringBuilder();
        Random random = new Random();
        for (int i = 0; i < ITEM_SIZE * 1024; i++) {
            sb.append(random.nextInt(10));
        }
        buffer = sb.toString();
        ExecutorService es = Executors.newCachedThreadPool();
        final CountDownLatch activeThreadCounter = new CountDownLatch(QUERY_THREAD_NUM + INSERT_THREAD_NUM);
        final Mongo mongo = new Mongo(HOST);
        mongo.setWriteConcern(WriteConcern.SAFE);
        try {
            for (int i = 0; i < INSERT_THREAD_NUM; i++) {
                es.execute(new Runnable() {

                    public void run() {

                        try {
                            DB db = mongo.getDB("test");
                            DBCollection collection = db.getCollection("test");
                            while (true) {
                                BasicDBObject doc = new BasicDBObject();
                                long id = CURRENT_ID.getAndIncrement();
                                doc.put("_id", id);
                                doc.put("name", buffer);
                                while (true) {
                                    try {
                                        collection.insert(doc);
                                        LAST_SUCCESS_ID.getAndIncrement();
                                        break;
                                    } catch (Exception e) {
                                        LOGGER.log(Level.SEVERE, String.format("Insert record %s failed", id), e);
                                    }
                                }
                            }
                        } finally {
                            activeThreadCounter.countDown();
                        }
                    }
                });
            }
            for (int i = 0; i < QUERY_THREAD_NUM; i++) {
                es.execute(new Runnable() {

                    public void run() {
                        try {
                            DB db = mongo.getDB("test");
                            DBCollection collection = db.getCollection("test");
                            Random random = new Random();
                            while (true) {
                                long lastID = LAST_SUCCESS_ID.get();
                                if (lastID >= 0) {
                                    lastID++;
                                    if (QUERY_RANGE > 0) {
                                        lastID = lastID > QUERY_RANGE ? QUERY_RANGE : lastID;
                                    }
                                    BasicDBObject primaryKey = new BasicDBObject();
                                    long id = Math.abs(random.nextLong()) % lastID;
                                    primaryKey.put("_id", id);
                                    while (true) {
                                        try {
                                            DBObject result = collection.findOne(primaryKey);
                                            break;
                                        } catch (Exception e) {
                                            LOGGER.log(Level.SEVERE, String.format("Query record %s failed", id), e);
                                        }
                                    }
                                }
                            }
                        } finally {
                            activeThreadCounter.countDown();
                        }
                    }
                });
            }
            es.shutdown();
            activeThreadCounter.await();
        } finally {
            mongo.close();
        }
    }
}
